/**
 * @author Gao Jian
 * @email gaoj@njupt.edu.cn
 * @create date 2021-06-03 10:38:05
 * @desc 连接破碎的小块
 */

#include "defs.h"

__global__ void ain_regions(unsigned char *output, unsigned char *input,
                            unsigned char value, unsigned int nw,
                            unsigned int nh) {
  unsigned nz = nw * nh;
  unsigned increment = blockDim.x * gridDim.x;
  unsigned nits = nz / increment + 1;
  unsigned idx = blockDim.x * blockIdx.x + threadIdx.x;

  for (unsigned i = 0; i < nits; i++, idx += increment) {
    if (idx >= nz)
      break;

    if (input[idx] == value) {
      output[idx] = 1; // 标记符号 = index+1
    } else {
      output[idx] = 0; // 区域外标记为0
    }
  }
}

void ainarize_regions(unsigned char *data, unsigned char value, unsigned int nw,
                      unsigned int nh) {
  unsigned nz = nw * nh;
  unsigned char *d_input = NULL;
  unsigned char *d_output = NULL;
  cudaMalloc((void **)&d_input, nz * sizeof(unsigned char));
  cudaMalloc((void **)&d_output, nz * sizeof(unsigned char));

  cudaMemcpy(d_input, data, nz * sizeof(unsigned char), cudaMemcpyHostToDevice);

  ain_regions<<<NUM_BLOCKS, NUM_THREADS>>>(d_output, d_input, value, nw, nh);

  cudaMemcpy(data, d_output, nz * sizeof(unsigned char),
             cudaMemcpyDeviceToHost);
  cudaFree(d_output);
  cudaFree(d_input);
}

__global__ void bin_regions(unsigned char *output, unsigned char *input,
                            unsigned char value, unsigned int nw,
                            unsigned int nh) {
  unsigned nz = nw * nh;
  unsigned increment = blockDim.x * gridDim.x;
  unsigned nits = nz / increment + 1;
  unsigned idx = blockDim.x * blockIdx.x + threadIdx.x;

  for (unsigned i = 0; i < nits; i++, idx += increment) {
    if (idx >= nz)
      break;

    if (input[idx] >= value && input[idx] != 255) {
      output[idx] = 1; // 标记符号 = index+1
    } else {
      output[idx] = 0; // 区域外标记为0
    }
  }
}

void binarize_regions(unsigned char *data, unsigned char value, unsigned int nw,
                      unsigned int nh) {
  unsigned nz = nw * nh;
  unsigned char *d_input = NULL;
  unsigned char *d_output = NULL;
  cudaMalloc((void **)&d_input, nz * sizeof(unsigned char));
  cudaMalloc((void **)&d_output, nz * sizeof(unsigned char));

  cudaMemcpy(d_input, data, nz * sizeof(unsigned char), cudaMemcpyHostToDevice);

  bin_regions<<<NUM_BLOCKS, NUM_THREADS>>>(d_output, d_input, value, nw, nh);

  cudaMemcpy(data, d_output, nz * sizeof(unsigned char),
             cudaMemcpyDeviceToHost);
  cudaFree(d_output);
  cudaFree(d_input);
}

// 标记符号 = index+1
__global__ void init_regions(unsigned int *regions, unsigned char *data,
                             unsigned char value, unsigned int nw,
                             unsigned int nh) {
  unsigned nz = nw * nh;
  unsigned increment = blockDim.x * gridDim.x;
  unsigned nits = nz / increment + 1;
  unsigned idx = blockDim.x * blockIdx.x + threadIdx.x;

  for (unsigned i = 0; i < nits; i++, idx += increment) {
    if (idx >= nz)
      break;

    if (data[idx] == value) {
      regions[idx] = idx + 1; // 标记符号 = index+1
    } else {
      regions[idx] = 0; // 区域外标记为0
    }
  }
}

__global__ void flip_regions(unsigned char *active, unsigned int *regions,
                             unsigned char mode, unsigned int nw,
                             unsigned int nh) {
  unsigned nz[4] = {nh, nh, nw, nw};

  unsigned ni[4] = {nw, nw, nh, nh};
  unsigned ns[4] = {nw, nw, 1, 1};
  unsigned oi[4] = {0, nw - 1, 0, nw * (nh - 1)};
  unsigned dz[4] = {1, -1, nw, -nw};
  unsigned increment = blockDim.x * gridDim.x;

  if (mode < 4) {
    unsigned tnz = nz[mode], tni = ni[mode], toi = oi[mode], tdz = dz[mode],
             tns = ns[mode];
    unsigned nits = tnz / increment + 1;
    unsigned idx = blockDim.x * blockIdx.x + threadIdx.x;

    char ba = 0;
    for (unsigned i = 0; i < nits; i++, idx += increment) {
      if (idx >= tnz)
        break;

      for (unsigned id = toi + idx * tns, n = 1; n < tni; id += tdz, n++)

        if (regions[id] != 0 && regions[id + tdz] != 0 &&
            regions[id] < regions[id + tdz]) {
          regions[id + tdz] = regions[id];
          if (ba == 0)
            ba = 1;
        }
    }
    if (ba == 1 && *active == 0)
      *active = 1;
  }
}

__global__ void kjoinparts(unsigned char *output, unsigned char *input,
                           unsigned int *regions, unsigned char value,
                           unsigned char mode, unsigned int nw,
                           unsigned int nh) {
  unsigned nz = nw * nh;
  unsigned increment = blockDim.x * gridDim.x;
  unsigned nits = nz / increment + 1;
  unsigned idx = blockDim.x * blockIdx.x + threadIdx.x;

  for (unsigned i = 0; i < nits; i++, idx += increment) {
    if (idx >= nz)
      break;

    if (input[idx] == value) {
      output[idx] = value;
      continue;
    }

    int iw = idx % nw;
    int ih = idx / nw;

    if (mode & (J1_1 | J1_2 | J1_4)) {

      if (iw < 1 || ih < 1 || iw > nw - 2 || ih > nh - 2)
        continue;
      if (regions && !(mode & J1_8)) {
        char act = 0;
        if (mode & J1_1) {
          int it[4] = {idx - 1, idx + 1, idx - nw, idx + nw};
          if (regions[it[0]] != 0 && regions[it[1]] != 0 &&
              regions[it[0]] != regions[it[1]])
            act = 1;
          if (regions[it[2]] != 0 && regions[it[3]] != 0 &&
              regions[it[2]] != regions[it[3]])
            act = 1;
        }
        if (mode & J1_2) {
          int it[4] = {idx - 1, idx + 1, idx - nw, idx + nw};
          if (regions[it[0]] != 0 && regions[it[2]] == 0 &&
              regions[it[3]] == 0 &&
              ((regions[it[1] - nw] != 0 &&
                regions[it[0]] != regions[it[1] - nw]) ||
               (regions[it[1] + nw] != 0 &&
                regions[it[0]] != regions[it[1] + nw])))
            act = 1;
          if (regions[it[2]] != 0 && regions[it[0]] == 0 &&
              regions[it[1]] == 0 &&
              ((regions[it[3] - 1] != 0 &&
                regions[it[2]] != regions[it[3] - 1]) ||
               regions[it[3] + 1] != 0 && regions[it[2]] != regions[it[3] + 1]))
            act = 1;
          if (regions[it[1]] != 0 && regions[it[2]] == 0 &&
              regions[it[3]] == 0 &&
              ((regions[it[0] - nw] != 0 &&
                regions[it[1]] != regions[it[0] - nw]) ||
               (regions[it[0] + nw] != 0 &&
                regions[it[1]] != regions[it[0] + nw])))
            act = 1;
          if (regions[it[3]] != 0 && regions[it[0]] == 0 &&
              regions[it[1]] == 0 &&
              ((regions[it[2] - 1] != 0 &&
                regions[it[3]] != regions[it[2] - 1]) ||
               (regions[it[2] + 1] != 0 &&
                regions[it[3]] != regions[it[2] + 1])))
            act = 1;
        }
        if (mode & J1_4) {
          int it[4] = {idx - 1 - nw, idx + 1 + nw, idx - nw + 1, idx + nw - 1};
          if (regions[it[0]] != 0 && regions[it[1]] != 0 &&
              regions[it[0]] != regions[it[1]])
            act = 1;
          if (regions[it[2]] != 0 && regions[it[3]] != 0 &&
              regions[it[2]] != regions[it[3]])
            act = 1;
        }
        if (act && output[idx] != value) {
          output[idx] = value;
        }
      } else {
        char act = 0;
        if (mode & J1_1) {
          int it[4] = {idx - 1, idx + 1, idx - nw, idx + nw};
          if (input[it[0]] == value && input[it[1]] == value)
            act = 1;
          if (input[it[2]] == value && input[it[3]] == value)
            act = 1;
        }
        if (mode & J1_2) {
          int it[4] = {idx - 1, idx + 1, idx - nw, idx + nw};
          if (input[it[0]] == value && input[it[2]] != value &&
              input[it[3]] != value &&
              (input[it[1] - nw] == value || input[it[1] + nw] == value))
            act = 1;
          if (input[it[2]] == value && input[it[0]] != value &&
              input[it[1]] != value &&
              (input[it[3] - 1] == value || input[it[3] + 1] == value))
            act = 1;
          if (input[it[1]] == value && input[it[2]] != value &&
              input[it[3]] != value &&
              (input[it[0] - nw] == value || input[it[0] + nw] == value))
            act = 1;
          if (input[it[3]] == value && input[it[0]] != value &&
              input[it[1]] != value &&
              (input[it[2] - 1] == value || input[it[2] + 1] == value))
            act = 1;
        }
        if (mode & J1_4) {
          int it[4] = {idx - 1 - nw, idx + 1 + nw, idx - nw + 1, idx + nw - 1};
          if (input[it[0]] == value && input[it[1]] == value)
            act = 1;
          if (input[it[2]] == value && input[it[3]] == value)
            act = 1;
        }
        if (act && output[idx] != value) {
          output[idx] = value;
        }
      }
    }

    if (mode & (J2_1 | J2_2 | J2_4)) {

      { // 水平双像素
        if (iw < 1 || ih < 1 || iw > nw - 3 || ih > nh - 2 ||
            input[idx + 1] == value)
          continue;
        if (regions && !(mode & J2_8)) {
          char act = 0;
          if (mode & J2_1) {
            int it[4] = {idx - 1, idx + 2};
            if (regions[it[0]] != 0 && regions[it[1]] != 0 &&
                regions[it[0]] != regions[it[1]])
              act = 1;
          }
          if (mode & J2_2) {
            int it[4] = {idx - 1, idx + 2};
            if (regions[it[0]] != 0 &&
                ((regions[it[1] - nw] != 0 &&
                  regions[it[0]] != regions[it[1] - nw]) ||
                 (regions[it[1] + nw] != 0 &&
                  regions[it[0]] != regions[it[1] + nw])))
              act = 1;
            if (regions[it[1]] != 0 &&
                ((regions[it[0] - nw] != 0 &&
                  regions[it[1]] != regions[it[0] - nw]) ||
                 (regions[it[0] + nw] != 0 &&
                  regions[it[1]] != regions[it[0] + nw])))
              act = 1;
          }
          if (mode & J2_4) {
            int it[4] = {idx - 1 - nw, idx + 2 + nw, idx - nw + 2,
                         idx + nw - 1};
            if (regions[it[0]] != 0 && regions[it[1]] != 0 &&
                regions[it[0]] != regions[it[1]])
              act = 1;
            if (regions[it[2]] != 0 && regions[it[3]] != 0 &&
                regions[it[2]] != regions[it[3]])
              act = 1;
          }
          if (act) {
            output[idx] = value;
            output[idx + 1] = value;
          }
        } else {
          char act = 0;
          if (mode & J2_1) {
            int it[4] = {idx - 1, idx + 2};
            if (input[it[0]] == value && input[it[1]] == value)
              act = 1;
          }
          if (mode & J2_2) {
            int it[4] = {idx - 1, idx + 2};
            if (input[it[0]] == value &&
                ((input[it[1] - nw] == value && input[idx - nw] != value &&
                  input[idx + 1 - nw] != value) ||
                 (input[it[1] + nw] == value && input[idx + nw] != value &&
                  input[idx + 1 + nw] != value)))
              act = 1;
            if (input[it[1]] == value &&
                ((input[it[0] - nw] == value && input[idx - nw] != value &&
                  input[idx + 1 - nw] != value) ||
                 (input[it[0] + nw] == value && input[idx + nw] != value &&
                  input[idx + 1 + nw] != value)))
              act = 1;
          }
          if (mode & J2_4) {
            int it[4] = {idx - 1 - nw, idx + 2 + nw, idx - nw + 2,
                         idx + nw - 1};
            if (input[it[0]] == value && input[it[1]] == value)
              act = 1;
            if (input[it[2]] == value && input[it[3]] == value)
              act = 1;
          }
          if (act) {
            output[idx] = value;
            output[idx + 1] = value;
          }
        }
      }
      { // 竖直双像素
        if (iw < 1 || ih < 1 || iw > nw - 2 || ih > nh - 3 ||
            input[idx + nw] == value)
          continue;
        if (regions && !(mode & J2_8)) {
          char act = 0;
          if (mode & J2_1) {
            int it[4] = {idx - nw, idx + 2 * nw};
            if (regions[it[0]] != 0 && regions[it[1]] != 0 &&
                regions[it[0]] != regions[it[1]])
              act = 1;
          }
          if (mode & J2_2) {
            int it[4] = {idx - nw, idx + 2 * nw};
            if (regions[it[0]] != 0 &&
                ((regions[it[1] - 1] != 0 &&
                  regions[it[0]] != regions[it[1] - 1]) ||
                 (regions[it[1] + 1] != 0 &&
                  regions[it[0]] != regions[it[1] + 1])))
              act = 1;
            if (regions[it[1]] != 0 &&
                ((regions[it[0] - 1] != 0 &&
                  regions[it[1]] != regions[it[0] - 1]) ||
                 (regions[it[0] + 1] != 0 &&
                  regions[it[1]] != regions[it[0] + 1])))
              act = 1;
          }
          if (mode & J2_4) {
            int it[4] = {idx - 1 - nw, idx + 1 + 2 * nw, idx - nw + 1,
                         idx + 2 * nw - 1};
            if (regions[it[0]] != 0 && regions[it[1]] != 0 &&
                regions[it[0]] != regions[it[1]])
              act = 1;
            if (regions[it[2]] != 0 && regions[it[3]] != 0 &&
                regions[it[2]] != regions[it[3]])
              act = 1;
          }
          if (act) {
            output[idx] = value;
            output[idx + nw] = value;
          }
        } else {
          char act = 0;
          if (mode & J2_1) {
            int it[4] = {idx - nw, idx + 2 * nw};
            if (input[it[0]] == value && input[it[1]] == value)
              act = 1;
          }
          if (mode & J2_2) {
            int it[4] = {idx - nw, idx + 2 * nw};
            if (input[it[0]] == value &&
                ((input[it[1] - 1] == value && input[idx + nw - 1] != value &&
                  input[idx + nw - 1] != value) ||
                 (input[it[1] + 1] == value && input[idx + nw + 1] != value &&
                  input[idx + nw + 1] != value)))
              act = 1;
            if (input[it[1]] == value &&
                ((input[it[0] - 1] == value && input[idx - 1] != value &&
                  input[idx + nw - 1] != value) ||
                 (input[it[0] + 1] == value && input[idx + 1] != value &&
                  input[idx + nw + 1] != value)))
              act = 1;
          }
          if (mode & J2_4) {
            int it[4] = {idx - 1 - nw, idx + 1 + 2 * nw, idx - nw + 1,
                         idx + 2 * nw - 1};
            if (input[it[0]] == value && input[it[1]] == value)
              act = 1;
            if (input[it[2]] == value && input[it[3]] == value)
              act = 1;
          }
          if (act) {
            output[idx] = value;
            output[idx + nw] = value;
          }
        }
      }
    }
  }
}

int join_parts(unsigned char *output, unsigned char *input, unsigned char value,
               unsigned char mode, unsigned int jnit, unsigned int nw,
               unsigned int nh, const unsigned verbose = 0) {
  unsigned int nz = nw * nh;
  unsigned char *d_input = NULL;
  unsigned char *d_output = NULL;
  unsigned int *d_regions = NULL;
  unsigned char *tmp = NULL;

  cudaMalloc((void **)&d_input, nz * sizeof(unsigned char));
  cudaMalloc((void **)&d_output, nz * sizeof(unsigned char));

  cudaMemcpy(d_input, input, nz * sizeof(unsigned char),
             cudaMemcpyHostToDevice);

  tmp = d_input;

  for (char i = 0; i < jnit; i++) {

    if (!(mode & J1_8) || !(mode & J2_8)) {
      cudaMalloc((void **)&d_regions, nz * sizeof(unsigned int));
      init_regions<<<NUM_BLOCKS, NUM_THREADS>>>(d_regions, d_input, value, nw,
                                                nh);

      unsigned char active;
      unsigned nit = 0;
      do {
        active = 0;
        cudaMemcpy(d_output, &active, sizeof(unsigned char),
                   cudaMemcpyHostToDevice);
        flip_regions<<<NUM_BLOCKS, NUM_THREADS>>>(d_output, d_regions, 0, nw,
                                                  nh);
        flip_regions<<<NUM_BLOCKS, NUM_THREADS>>>(d_output, d_regions, 2, nw,
                                                  nh);
        flip_regions<<<NUM_BLOCKS, NUM_THREADS>>>(d_output, d_regions, 1, nw,
                                                  nh);
        flip_regions<<<NUM_BLOCKS, NUM_THREADS>>>(d_output, d_regions, 3, nw,
                                                  nh);
        cudaMemcpy(&active, d_output, sizeof(unsigned char),
                   cudaMemcpyDeviceToHost);
        nit++;
        // printf(" time[%4d] === >>>>>> \n", nit);
      } while (active);
    }

    cudaMemcpy(d_output, output, nz * sizeof(unsigned char),
               cudaMemcpyHostToDevice);
    kjoinparts<<<NUM_BLOCKS, NUM_THREADS>>>(d_output, d_input, d_regions, value,
                                            mode, nw, nh);

    tmp = d_output;
    d_output = d_input;
    d_input = tmp;

    if (verbose)
      printf("+");
    if ((i + 1) % 50 == 0 && verbose)
      printf("  [%3d]\n  ", i + 1);
  }
  if (verbose)
    printf("  [%3d]\n", jnit);

  // cudaMemcpy(temp, d_regions, nz * sizeof(unsigned int),
  // cudaMemcpyDeviceToHost);
  cudaMemcpy(output, tmp, nz * sizeof(unsigned char), cudaMemcpyDeviceToHost);

  if (!(mode & J1_8) || !(mode & J2_8)) {
    cudaFree(d_regions);
  }
  cudaFree(d_output);
  cudaFree(d_input);

  return 0;
}

void PrintDeviceInfo() {
  int deviceCount = 0;
  int devID = 0;
  int curID = 0;

  struct cudaDeviceProp deviceProp;
  cudaGetDevice(&curID);

  cudaGetDeviceCount(&deviceCount);

  for (devID = 0; devID < deviceCount; devID++) {

    cudaGetDeviceProperties(&deviceProp, devID);

    if (deviceProp.computeMode == cudaComputeModeProhibited) {
      fprintf(stderr,
              "Error: device %d in <Compute Mode Prohibited>,"
              " no threads can use ::cudaSetDevice().\n",
              devID);
      return;
    }

    if (devID == curID)
      printf("  ***>>> GPU Device %d: \"%s\" with compute capability %d.%d\n",
             devID, deviceProp.name, deviceProp.major, deviceProp.minor);
    else
      printf("  +++>>> GPU Device %d: \"%s\" with compute capability %d.%d\n",
             devID, deviceProp.name, deviceProp.major, deviceProp.minor);
  }

  printf("\n");
}
